// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.
{% include "utils/base.groovy.j2" with context %}
{% import 'utils/macros.j2' as m with context -%}

{% call m.invoke_build(
  name='BUILD: GPU',
  node='CPU-SMALL',
  condition='!skip_ci',
  ws='tvm/build-gpu',
  docker_image='ci_gpu',
) %}
  sh "${docker_run} --no-gpu ${ci_gpu} ./tests/scripts/task_config_build_gpu.sh build"
  cmake_build("${ci_gpu} --no-gpu", 'build', '-j2')
  make_standalone_crt("${ci_gpu} --no-gpu", 'build')
  {{ m.upload_artifacts(tag='gpu', filenames=tvm_multilib + tvm_allvisible + microtvm_template_projects + crttest + standalone_crt) }}

  // compiler test
  sh "rm -rf build"
  sh "${docker_run} --no-gpu ${ci_gpu} ./tests/scripts/task_config_build_gpu_other.sh build"
  cmake_build("${ci_gpu} --no-gpu", 'build', '-j2')
  make_standalone_crt("${ci_gpu} --no-gpu", 'build')
  {{ m.upload_artifacts(tag='gpu2', filenames=tvm_lib + crttest + standalone_crt) }}
{% endcall %}

{% set test_method_names = [] %}

{% call(shard_index, num_shards) m.sharded_test_step(
  name="unittest: GPU",
  num_shards=3,
  node="GPU",
  ws="tvm/ut-python-gpu",
  platform="gpu",
  docker_image="ci_gpu",
  test_method_names=test_method_names,
) %}
  {% if shard_index == 1 %}
  {{ m.download_artifacts(tag='gpu2') }}
  sh "${docker_run} --no-gpu ${ci_gpu} ./tests/scripts/task_config_build_gpu_other.sh build"
  // These require a GPU to finish the build (i.e. CUDA needs to be load-able)
  make_standalone_crt(ci_gpu, 'build')
  // make_cpp_tests(ci_gpu, 'build')
  // cpp_unittest(ci_gpu)

  sh "rm -rf build"
  {{ m.download_artifacts(tag='gpu') }}
  ci_setup(ci_gpu)
  sh "${docker_run} --no-gpu ${ci_gpu} ./tests/scripts/task_config_build_gpu.sh build"
  make_standalone_crt(ci_gpu, 'build')
  make_cpp_tests(ci_gpu, 'build')
  cpp_unittest(ci_gpu)
  sh (
    script: "${docker_run} ${ci_gpu} python3 ./tests/scripts/task_build.py --sccache-bucket tvm-sccache-prod --cmake-target opencl-cpptest --build-dir build",
    label: 'Make OpenCL cpp unit tests',
  )
  sh (
    script: "${docker_run} ${ci_gpu} ./tests/scripts/task_opencl_cpp_unittest.sh",
    label: 'Run OpenCL cpp unit tests',
  )
  micro_cpp_unittest(ci_gpu)
  {% else %}
  {{ m.download_artifacts(tag='gpu') }}
  ci_setup(ci_gpu)
  {% endif %}
  {% if shard_index == 2 or num_shards < 2 %}
  sh (
    script: "${docker_run} ${ci_gpu} ./tests/scripts/task_java_unittest.sh",
    label: 'Run Java unit tests',
  )
  {% endif %}
  sh (
    script: "${docker_run} ${ci_gpu} ./tests/scripts/task_python_unittest_gpuonly.sh",
    label: 'Run Python GPU unit tests',
  )
  sh (
    script: "${docker_run} ${ci_gpu} ./tests/scripts/task_python_integration_gpuonly.sh",
    label: 'Run Python GPU integration tests',
  )
{% endcall %}

{% call(shard_index, num_shards) m.sharded_test_step(
  name="topi: GPU",
  node="GPU",
  num_shards=3,
  ws="tvm/topi-python-gpu",
  platform="gpu",
  docker_image="ci_gpu",
  test_method_names=test_method_names,
) %}
  {{ m.download_artifacts(tag='gpu') }}
  ci_setup(ci_gpu)
  sh (
    script: "${docker_run} ${ci_gpu} ./tests/scripts/task_python_topi.sh",
    label: 'Run TOPI tests',
  )
{% endcall %}

{% call(shard_index, num_shards) m.sharded_test_step(
  name="frontend: GPU",
  node="GPU",
  num_shards=6,
  ws="tvm/frontend-python-gpu",
  platform="gpu",
  docker_image="ci_gpu",
  test_method_names=test_method_names,
) %}
  {{ m.download_artifacts(tag='gpu') }}
  ci_setup(ci_gpu)
  sh (
    script: "${docker_run} ${ci_gpu} ./tests/scripts/task_python_frontend.sh",
    label: 'Run Python frontend tests',
  )
{% endcall %}

{% call(shard_index, num_shards) m.sharded_test_step(
  name="docs: GPU",
  node="GPU",
  num_shards=1,
  ws="tvm/docs-python-gpu",
  platform="gpu",
  docker_image="ci_gpu",
  condition="!skip_ci",
  test_method_names=test_method_names,
) %}
  {{ m.download_artifacts(tag='gpu') }}
  ci_setup(ci_gpu)
  sh (
    script: "${docker_run} ${ci_gpu} ./tests/scripts/task_python_docs.sh",
    label: 'Build docs',
  )
  {{ m.upload_artifacts(tag='docs', filenames=["docs.tgz"]) }}
  sh(
    script: "aws s3 cp --no-progress _docs s3://${s3_bucket}/${s3_prefix}/docs --recursive",
    label: 'Upload docs to S3',
  )
{% endcall %}


{{ m.invoke_tests(test_method_names) }}


def deploy_docs() {
  // Note: This code must stay in the Jenkinsfile to ensure that it runs
  // from a trusted context only
  sh(
    script: '''
      set -eux
      rm -rf tvm-site
      git clone -b $DOCS_DEPLOY_BRANCH --depth=1 https://github.com/apache/tvm-site
      cd tvm-site
      git status
      git checkout -B $DOCS_DEPLOY_BRANCH

      git ls-tree HEAD docs/ --name-only | grep -vP '^docs/v\\d' | xargs rm -rf
      mkdir -p docs
      tar xf ../docs.tgz -C docs
      COMMIT=$(cat docs/commit_hash)
      git add .
      git config user.name tvm-bot
      git config user.email 95660001+tvm-bot@users.noreply.github.com
      git commit -m"deploying docs (apache/tvm@$COMMIT)"
      git status
    ''',
    label: 'Unpack docs and update tvm-site'
  )

  withCredentials([string(
    credentialsId: 'docs-push-token',
    variable: 'GITHUB_TOKEN',
    )]) {
    sh(
      script: '''
        cd tvm-site
        git remote add deploy https://$GITHUB_TOKEN:x-oauth-basic@github.com/apache/tvm-site.git
        git push deploy $DOCS_DEPLOY_BRANCH || true
      ''',
      label: 'Upload docs to apache/tvm-site'
    )
  }
}

def deploy() {
  stage('Deploy') {
    if (env.BRANCH_NAME == 'main') {
      parallel(
        {% call m.deploy_step(
          name="Deploy Docs",
          feature_flag="env.DOCS_DEPLOY_ENABLED == 'yes'",
          ws="tvm/deploy-docs",
        ) %}
          init_git()
          sh(
            script: "./${jenkins_scripts_root}/s3.py --action download --bucket ${s3_bucket} --prefix ${s3_prefix}/docs --items docs.tgz",
            label: 'Download docs folder from S3',
          )
          deploy_docs()
        {% endcall %}
      )
    }
  }
}

deploy()
